Step Functionsで扱えるSageMakerのAPIが増えました

どうも、DA部の大澤です。

Step FunctionsのSageMakerとの統合機能がアップデートされました！

これまでは学習ジョブの作成とバッチ変換ジョブの作成にのみ対応していたStep FunctionsのSageMakerとの統合機能ですが、今回のアップデートによりモデルの作成やハイパーパラメータチューニングジョブの作成、ラベリングジョブの作成などなど新たに多くのAPIへの対応がなされました。

• AWS Step Functions expands Amazon SageMaker service integration

対応API

2019年10月4日現在、Step Functionsが対応しているSageMakerのAPIは次の通りです。上2つがもともと対応されていたAPIで残りが今回新たに対応されたものとなります。

• CreateTrainingJob
• CreateTransformJob
• CreateEndpoint
• CreateEndpointConfig
• CreateHyperParameterTuningJob
• CreateLabelingJob
• CreateModel
• UpdateEndpoint

詳細についてドキュメントをご覧ください。

• Manage Amazon SageMaker with Step Functions - AWS Step Functions

やってみる

今回のアップデートでは対応APIの拡張とともに2つのサンプルプロジェクトの追加も行われました。"機会学習モデルをトレーニング"と"機械学習モデルをチューニング"という内容です。

その内の"機械学習モデルをトレーニング”を使ってみたので、その内容を紹介します。

ステートマシーン

実行するステートマシーンは次のようなものとなります。

1. Lambdaでデータセットを作成する
   • trainとtestの2つのデータを作成
   • trainのデータ形式: f(x, y), x, y
   • testのデータ形式: x, y
   • xとyは0から100までのランダムな整数
   • f(x, y)はy >= 3x + 5を満たす場合は1、満たさない場合は0を返す関数
2. SageMakerでXGBoostの回帰モデルをtrainデータを使って学習させる
3. 2.で得られたモデルアーティファクトからSageMakerでモデルを作成する
4. 3.で作成したモデルを使って、SageMakerのバッチ変換ジョブを作成し、testデータを推論する

ステートマシーンの定義は次の通りです。

{
  "StartAt": "Generate dataset",
  "States": {
    "Generate dataset": {
      "Resource": "<GENERATE_LAMBDA_FUNCTION_ARN>",
      "Type": "Task",
      "Next": "Train model (XGBoost)"
    },
    "Train model (XGBoost)": {
      "Resource": "arn:<PARTITION>:states:::sagemaker:createTrainingJob.sync",
      "Parameters": {
        "AlgorithmSpecification": {
          "TrainingImage": "<SAGEMAKER_TRAINING_IMAGE>",
          "TrainingInputMode": "File"
        },
        "OutputDataConfig": {
          "S3OutputPath": "s3://<S3_BUCKET>/models"
        },
        "StoppingCondition": {
          "MaxRuntimeInSeconds": 86400
        },
        "ResourceConfig": {
          "InstanceCount": 1,
          "InstanceType": "ml.m4.xlarge",
          "VolumeSizeInGB": 30
        },
        "RoleArn": "<SAGEMAKER_EXECUTION_ROLE_ARN>",
        "InputDataConfig": [
          {
            "DataSource": {
              "S3DataSource": {
                "S3DataDistributionType": "ShardedByS3Key",
                "S3DataType": "S3Prefix",
                "S3Uri": "s3://<S3_BUCKET>/csv/train.csv"
              }
            },
            "ChannelName": "train",
            "ContentType": "text/csv"
          }
        ],
        "HyperParameters": {
          "objective": "reg:logistic",
          "eval_metric": "rmse",
          "num_round": "5"
        },
        "TrainingJobName.$": "$$.Execution.Name"
      },
      "Type": "Task",
      "Next": "Save Model"
    },
    "Save Model": {
      "Parameters": {
        "PrimaryContainer": {
          "Image": "<SAGEMAKER_TRAINING_IMAGE>",
          "Environment": {},
          "ModelDataUrl.$": "$.ModelArtifacts.S3ModelArtifacts"
        },
        "ExecutionRoleArn": "<SAGEMAKER_EXECUTION_ROLE_ARN>",
        "ModelName.$": "$.TrainingJobName"
      },
      "Resource": "arn:aws:states:::sagemaker:createModel",
      "Type": "Task",
      "Next": "Batch transform"
    },
    "Batch transform": {
      "Type": "Task",
      "Resource": "arn:<PARTITION>:states:::sagemaker:createTransformJob.sync",
      "Parameters": {
        "ModelName.$": "$$.Execution.Name",
        "TransformInput": {
          "CompressionType": "None",
          "ContentType": "text/csv",
          "DataSource": {
            "S3DataSource": {
              "S3DataType": "S3Prefix",
              "S3Uri": "s3://<S3_BUCKET>/csv/test.csv"
            }
          }
        },
        "TransformOutput": {
          "S3OutputPath": "s3://<S3_BUCKET>/output"
        },
        "TransformResources": {
          "InstanceCount": 1,
          "InstanceType": "ml.m4.xlarge"
        },
        "TransformJobName.$": "$$.Execution.Name"
      },
      "End": true
    }
  }
}

ステートマシーンの作成

マネジメントコンソールでStep Functionsのステートマシーン一覧の画面を開き、画面右上のボタンからステートマシーンを作成画面を開きます。 サンプルプロジェクトの一覧から"機械学習モデルをトレーニング"を選び、"次へ"で確認画面へ進みます。

内容を確認し、"リソースのデプロイ"で必要なリソースのデプロイを実施します。 これによりCloudFormationで一通り必要なリソースが自動で作成されます。

ステートマシーンの実行

リソースのデプロイが完了すると、次のような画面が自動的に開きます。"実行の開始"をクリックし、ステートマシーンを実行します。

10分弱程で実行が完了し、実行のステータスが完了へと変化します。

SageMakerのトレーニングジョブ、モデル、バッチ変換ジョブはステートマシンの実行IDと同じ名前で作成されています。

Step Functionsのステートマシンによって、回帰モデルを学習し、バッチ変換ジョブが実行されていることが確認できました。

リソースの削除

確認が終わり今回使用したリソースが不要になったので、CloudFormationのスタックを削除します。

さいごに

今回はStep FunctionsのSageMakerとの統合機能のアップデートとサンプルプロジェクトについて紹介いたしました。
今回のアップデートによりこれまではLambda関数で処理を書かざるを得なかったものも、SageMakerのタスクとしてステートマシーンに定義することができるようになります。うまく活用することで機械学習システムの構築、運用の負荷軽減に繋がりそうです。